# -*- coding: utf-8 -*-
'''
@file: utils.py
@time: 19-1-11 上午10:24
@desc:

'''
from __future__ import division
from __future__ import division

import glob

import imgaug as ia
from  imgaug import augmenters as iaa
import random
import os
import sys

import cv2
import numpy as np
from datetime import datetime

import pandas as pd
import matplotlib.pyplot as plt
from skimage.util import random_noise

import SimpleITK as sitk


def mem_info():
    import subprocess
    dev = subprocess.check_output(
        "nvidia-smi | grep MiB | awk -F '|' '{print $3}' | awk -F '/' '{print $1}' | grep -Eo '[0-9]{1,10}'",
        shell=True)
    dev = dev.decode()
    dev_mem = list(map(lambda x: int(x), dev.split('\n')[:-1]))
    return dev_mem


def add_pypath(path):
    if path not in sys.path:
        sys.path.insert(0, path)


def make_link(dest_path, link_path):
    os.system('ln -s {} {}'.format(dest_path, link_path))


def make_dir(path):
    if os.path.exists(path) or os.path.islink(path):
        return
    os.makedirs(path)
    print("make dir: {}".format(path))


def exists(dir):
    if not os.path.exists(dir):
        raise IOError("{} not exist!".format(dir))


def del_file(path, msg='{} deleted.'):
    if os.path.exists(path):
        os.remove(path)
        print(msg.format(path))


def approx_equal(a, b, eps=1e-9):
    return np.fabs(a - b) < eps


def get_rng(obj=None):
    """
    Get a good RNG seeded with time, pid and the object.

    Args:
        obj: some object to use to generate random seed.
    Returns:
        np.random.RandomState: the RNG.
    """
    seed = (id(obj) + os.getpid() +
            int(datetime.now().strftime("%Y%m%d%H%M%S%f"))) % 4294967295
    return np.random.RandomState(seed)



def get_average_value():
    """计算平均值的闭包。
    a = get_average_value()
    a(2)
    Out[11]: 2.0
    """
    
    s = [0.]  ## sum
    count = [0.]
    
    def func(value):
        s[0] += value
        count[0] += 1
        return s[0] / count[0]
    
    return func


def add_gaussian_noise(arr, mean=0, var=0.01):
    ## 可以是3d数据
    return random_noise(arr, mode='gaussian', seed=2333, clip=True, mean=mean, var=var)


def min_max_norm(image, multi_255=True):
    """min-max normal"""
    maxval = np.max(image)
    minval = np.min(image)
    # print 'max {}, min {}'.format(maxval, minval)
    image = (image - minval) / float(maxval - minval)
    if multi_255:
        image = image * 255
    return image.astype(np.float32)


def read_txt(filepath):
    """读取txt中 uid用"""
    if not isinstance(filepath, str):
        raise ValueError("input error, filepath must be string type: %s" % type(filepath))
    
    if not os.path.exists(filepath):
        raise ValueError("%s file or path not exist!" % filepath)
    
    f = open(filepath)
    lists = f.read().splitlines()
    f.close()
    lists = [uid for uid in lists if len(uid) != 0]
    
    return lists


#######################
###   make heatmap
#########################

def array2keypoint(coords, image_shape):
    keypoints_list = [ia.Keypoint(x=x, y=y) for x, y in coords]
    keypoints = ia.KeypointsOnImage(keypoints_list, shape=image_shape)
    return keypoints


def image_keypoints_aug_rotate(image, coords):
    """ image 和 keypoint一起做aug。F

    :param image:
    :param coords: array, [x, y]
    :return:
    """
    # ia.seed(2333)
    
    keypoints_list = [ia.Keypoint(x=x, y=y) for x, y in coords]
    keypoints = ia.KeypointsOnImage(keypoints_list, shape=image.shape)
    
    seq = iaa.Sequential([
        iaa.Affine(
            rotate=random.randint(-35, 35),
            # scale=(0.7, 1.1)
        ),
    
    ])
    
    # for _ in range(100):
    seq_det = seq.to_deterministic()
    image_aug = seq_det.augment_images([image])[0]
    keypoints_aug = seq_det.augment_keypoints([keypoints])[0]
    # print keypoints_aug
    # image_before = keypoints.draw_on_image(image, size=3)
    # image_after = keypoints_aug.draw_on_image(image_aug, size=3)
    coords_aug = keypoints_aug.get_coords_array()
    
    return image_aug, coords_aug


def image_keypoints_aug_zoom(image, coords):
    """ image 和 keypoint一起做aug。F

    :param image:
    :param coords: array, [x, y]
    :return:
    """
    # ia.seed(2333)
    
    keypoints_list = [ia.Keypoint(x=x, y=y) for x, y in coords]
    keypoints = ia.KeypointsOnImage(keypoints_list, shape=image.shape)
    
    seq = iaa.Sequential([
        iaa.Affine(
            
            scale=(0.8, 1.1)
        ),
    
    ])
    
    # for _ in range(100):
    seq_det = seq.to_deterministic()
    image_aug = seq_det.augment_images([image])[0]
    keypoints_aug = seq_det.augment_keypoints([keypoints])[0]
    # print keypoints_aug
    # image_before = keypoints.draw_on_image(image, size=3)
    # image_after = keypoints_aug.draw_on_image(image_aug, size=3)
    coords_aug = keypoints_aug.get_coords_array()
    
    return image_aug, coords_aug


def image_keypoints_aug_flip(image, coords):
    """ image 和 keypoint一起做aug。F

    :param image:
    :param coords: array, [x, y]
    :return:
    """
    # ia.seed(2333)
    
    keypoints_list = [ia.Keypoint(x=x, y=y) for x, y in coords]
    keypoints = ia.KeypointsOnImage(keypoints_list, shape=image.shape)
    
    seq = iaa.Sequential([
        
        iaa.Fliplr(0.5)
    
    ])
    
    seq_det = seq.to_deterministic()
    image_aug = seq_det.augment_images([image])[0]
    keypoints_aug = seq_det.augment_keypoints([keypoints])[0]
    
    coords_aug = keypoints_aug.get_coords_array()
    
    return image_aug, coords_aug


def image_keypoints_aug_kuanguanjie(image, coords):
    """ image 和 keypoint一起做aug。F

    :param image:
    :param coords: array, [x, y]
    :return:
    """
    # ia.seed(2333)
    
    keypoints_list = [ia.Keypoint(x=x, y=y) for x, y in coords]
    keypoints = ia.KeypointsOnImage(keypoints_list, shape=image.shape)
    
    seq = iaa.Sequential([
        iaa.Affine(
            rotate=random.randint(-20, 20),
            scale=(0.75, 1.3)
        )
    ])
    
    seq_det = seq.to_deterministic()
    image_aug = seq_det.augment_images([image])[0]
    keypoints_aug = seq_det.augment_keypoints([keypoints])[0]
    
    coords_aug = keypoints_aug.get_coords_array()
    
    return image_aug, coords_aug


def make_heatmap(coords, target_image_size, src_image_size, gaussion_kernel, keypoints_num):
    """ 构建热图。k层。

    :param coords: array, [x, y]， 原始图像的坐标,即 width, hight
    :param target_image_size: 输出热图的大小
    :param src_image_size:  原始图像的大小, 热图的坐标点的原来大小
    :param gaussion_kernel: 高斯核的大小
    :param keypoints_num: 关键点的个数，所有。
    :return: heatmap array，返回k个关键点的，k层热图，一层一个关键点。
    """
    ## 输出的热图， batchsize， keypoints num, image size
    hm_output = np.zeros([keypoints_num, target_image_size[0], target_image_size[1]], dtype=np.float32)
    for j in range(keypoints_num):
        # print 'j: ',j
        if ((coords[j, 0] < 0) or (coords[j, 1] < 0)
        
            or (coords[j, 0] > src_image_size[1])
            or (coords[j, 1] > src_image_size[0])):
            ## 如果有坐标小于1或超过边界那说明关键点不在图像里面，或许是aug的时候到外面去了，也可能是本身就是不可见.
            ## 上面的判断先和原图比较，因为坐标还没做缩放
            continue
        # print j
        # 给热图赋值
        x_tmp = int(round(coords[j, 0] * (target_image_size[1] / float(src_image_size[1]))))
        y_tmp = int(round(coords[j, 1] * (target_image_size[0] / float(src_image_size[0]))))
        
        if (x_tmp >= target_image_size[1]) or (y_tmp >= target_image_size[0]):
            ## 如果缩放后行或列超过边界那么就不要，其实是四舍五入造成的。
            continue
        
        hm_output[j, y_tmp, x_tmp] = 1
        ## 高斯变换,直接在这里做高斯，这样没有特征点的层就不需要再做滤波了。
        hm_output[j] = cv2.GaussianBlur(hm_output[j], ksize=gaussion_kernel, sigmaX=0, sigmaY=0)
        ## 归一化到 0-1之间
        am = np.amax(hm_output[j])  # 放回最大值
        hm_output[j] = hm_output[j] / float(am)
    hm_output *= 255  ## ???
    
    return hm_output




def get_uid_from_path(path):
    return os.path.splitext(os.path.split(path)[1])[0]




#######
def heatmap_on_image(image, heatmap, save=None, title='', alpha=0.5, display=False, axis='on',
                     verbose=False):
    ## normalize heat map
    maxval = np.max(heatmap)
    minval = np.min(heatmap)
    heatmap = (heatmap - minval) / float(maxval - minval)
    heatmap = np.uint8(heatmap * 255)
    
    ## display
    im1 = plt.imshow(image, cmap='gray')
    im2 = plt.imshow(heatmap, cmap='viridis', alpha=alpha)
    if title:
        plt.title(title)
    
    plt.axis(axis)
    
    if display:
        plt.show()
    
    if save is not None:
        if verbose:
            print("save image: {}".format(save))
        plt.savefig(save, bbox_inches='tight', pad_inches=0)


def monochrome_convert(image, default_window=True):
    # pminterp = image.header.metadata.get('0028|0004')
    ## 0020|0010 study id 0020|000d study instance id
    # print 'study ID: {}; study instance ID: {}'.format(image.GetMetaData('0020|0010'), image.GetMetaData('0020|000d'))
    pminterp = image.GetMetaData('0028|0004')
    if pminterp:
        pminterp = pminterp.strip().upper()
    
    # print 'pminterp: ', pminterp
    # pixels = image.pixel_data
    pixels = sitk.GetArrayFromImage(image)
    # print "np.max(pixels), np.min(pixels): ", np.max(pixels), np.min(pixels)
    
    if pminterp == 'MONOCHROME1':
        bits_stored = image.GetMetaData('0028|0101')
        # print bits_stored
        if bits_stored:
            maxval = 2 ** int(bits_stored) - 1
        else:
            maxval = 4095  # A sane default (12-bit scanner)
        
        if default_window:
            window_center = image.GetMetaData("0028|1050")
            window_width = image.GetMetaData("0028|1051")
            if window_center and window_width:
                # print 'maxval: ', maxval
                window_center = float(window_center)
                window_width = float(window_width)
                # print window_center, type(window_center), window_width, type(window_width)
                # print "(int(window_center) + int(window_width) / 2): ",(int(window_center) + int(window_width) / 2)
                # print "(int(window_center) - int(window_width) / 2): ",(int(window_center) - int(window_width) / 2)
                pixels[pixels > (int(window_center) + int(window_width) / 2)] = int(window_center) + int(
                    window_width) / 2
                pixels[pixels < (int(window_center) - int(window_width) / 2)] = int(window_center) - int(
                    window_width) / 2
                # print ' np.max(pixels), np.min(pixels)',np.max(pixels), np.min(pixels)
        # print 'maxval: ', maxval
        
        pixels = maxval - pixels
        # print pixels.dtype
        # print np.max(pixels), np.min(pixels)
    
    elif pminterp == 'MONOCHROME2':
        if default_window:
            window_center = image.GetMetaData("0028|1050")
            window_width = image.GetMetaData("0028|1051")
            if window_center and window_width:
                pixels[pixels > (int(window_center) + int(window_width) / 2)] = int(window_center) + int(
                    window_width) / 2
                pixels[pixels < (int(window_center) - int(window_width) / 2)] = int(window_center) - int(
                    window_width) / 2
    
    elif pminterp == 'RGB':
        colorweights = [0.299, 0.587, 0.114]
        pixels = np.mean(pixels * colorweights, axis=2)
    
    return pixels


def draw_circle_label(image_src, coords_tmp, coord_image_shape, tag_list, thickness=2, fontScale=1.4, radius=5,
                      text=True, color=(255, 0, 0)):
    """

    :param image: image
    :param coords: 坐标 [x, y]
    :param coord_image_shape: 坐标所在的图像大小 (hight, width)
    :param tag_list: 每个坐标对应的标签 , list
    :return: 标注好的图像
    """
    coords = coords_tmp.copy()
    assert coords.shape[0] == len(tag_list), "coord number({}) not equal tag number({}).".format(coords.shape[0],
                                                                                                 len(tag_list))
    image = image_src.copy()
    zoom_rvs = (image_src.shape[0] / coord_image_shape[0], image_src.shape[1] / coord_image_shape[1])
    coords[..., 0] = np.round(coords[..., 0] * zoom_rvs[1]).astype(np.uint16)
    coords[..., 1] = np.round(coords[..., 1] * zoom_rvs[0]).astype(np.uint16)
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    for k in range(coords.shape[0]):
        center = tuple(coords[k])  ## x,y
        if center[0] < 0 or center[1] < 0:
            continue
        cv2.circle(image, center=center, radius=radius, color=color, thickness=-1)
        if text:
            cv2.putText(image, text=tag_list[k], org=(center[0] + 9, center[1] + 5),
                        fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=fontScale,
                        color=color, thickness=thickness)
    return image


def heatmap2coord(heatmap, if_gaussion=True):
    """ 2d """
    ## 高斯平滑
    # print np.amax(heatmap)
    heatmap /= np.amax(heatmap)
    # print np.amax(heatmap)
    if if_gaussion:
        heatmap = cv2.GaussianBlur(heatmap, (21, 21), 0)
    ### get max value coord
    # print heatmap.shape
    
    max_val = heatmap.argmax()
    # print 'max val: ', max_val
    y, x = np.unravel_index(max_val, heatmap.shape)
    return (x, y)


def get_valid(coords):
    """ 获得哪个点是存在（有效）的，这样计算loss的时候不计算不存在的点.
            有效的关键点，即加入计算loss的关键点。
    :param coords: array , [[x1, y1] [x2, y2] .....]
    """
    valid = coords > 0
    valid = np.float32(np.logical_and(valid[:, 0], valid[:, 1]))
    return valid


def read_dicom2npz(dicom_main_dir, csv, save_dir):
    ## get uid and dicom
    df = pd.read_csv(csv)
    uid_list = set(df['series_uid'])
    dicomfile_list = []  ## dicom file
    # print uid_list
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    
    print
    "uid length: ", len(uid_list)
    for i, uid in enumerate(uid_list):
        print
        'uid', i, uid
        dicomfile = glob.glob(os.path.join(dicom_main_dir, "*/%s/*dcm" % uid))
        # print dicomfile
        try:
            dicomfile_list.append(dicomfile[0])
        except:
            print
            '{} not exist in {}'.format(uid, dicom_main_dir)
            continue
            # print dicomfile[0]
    print
    'dicom length: ', len(dicomfile_list)
    
    ## read dicom to npz
    for i, dicomfile in enumerate(dicomfile_list):
        print
        i
        uid = os.path.split(os.path.split(dicomfile)[0])[1]
        print
        uid
        print
        dicomfile
        # if os.path.exists(os.path.join(save_dir, "%s_image.npz" % uid)):
        #     continue
        scan = sitk.ReadImage(dicomfile)
        print
        'monochrome_convert '
        image_tmp = monochrome_convert(scan)
        print
        'monochrome_convert ok'
        st_id = scan.GetMetaData('0020|0010')
        
        spacing = scan.GetSpacing()
        print
        'spacing: ', spacing
        print
        image_tmp.dtype
        if save_dir is not None:
            np.savez_compressed(os.path.join(save_dir, "{}_{}.npz".format(st_id, uid)), image=image_tmp)


def plot(*images):
    for i, image in enumerate(images, start=1):
        plt.subplot(1, len(images), i)
        plt.imshow(image, cmap='gray')
    plt.show()


class writeCSV():
    """追加写入csv。

     example:
     csv_savepath = '/home/xupan/Tb/keypoint_detection/kuan_guan_jie/predict/dataset2/val_10point_EPOCH30/measure_predict.csv'
    csv_column = ("series_uid", "lateral_ce_angle", "tonnis_angle", "sharp_angle", "femoral_head_extrusion_index")
    write_csv = utils.writeCSV(csv_savepath, csv_column)
    write_csv.make_csv()
    for循环写入
        write_csv.append_write_csv({"series_uid": index,
                                    "lateral_ce_angle":m.lateral_ce_angle ,
                                    "tonnis_angle": m.tonnis_angle,
                                    "sharp_angle": m.sharp_angle,
                                    "femoral_head_extrusion_index": m.femoral_head_extrusion_index
                                    })
    """
    
    def __init__(self, csv_save_path, df_col, index=False):
        self.csv_save_path = csv_save_path
        self.df_col = df_col
        self.index = index
        # if not os.path.exists(csv_save_path):
        #     self.make_csv()
        self.make_csv()
    
    def make_csv(self):
        # if os.path.exists(self.csv_save_path):
        #     raise IOError("{} exist! please theck.".format(self.csv_save_path))
        df = pd.DataFrame(columns=self.df_col)
        df.to_csv(self.csv_save_path, mode='w', index=self.index)
    
    def append_write_csv(self, value_dict):
        df = pd.DataFrame(columns=self.df_col)
        df = df.append(value_dict, ignore_index=True)
        df.to_csv(self.csv_save_path, mode='a+', index=self.index, header=False)


if __name__ == '__main__':
    # aug_and_heatmap_test()
    checkpoint_variable_rename(
        checkpoint_dir='/home/xupan/Tb/keypoint_detection/checkpoint_all/pre_train_model/resnet_v1_50.ckpt',
        replace_from=None,
        replace_to=None,
        add_prefix="stage4/",
        checkpoint_save_path='/home/xupan/Tb/keypoint_detection/checkpoint_all/pre_train_model/stage4_resnet_v1_50.ckpt',
        dry_run=False)
